#  -*- coding: utf-8 -*-

from datetime import datetime

from pandas import DataFrame
from pymongo import ASCENDING
import time
from util.database import DB_CONN
from factor.change_rate_factor import ChangeRateFactor
from factor.pe_factor import PEFactor
from factor.hfq_ma_factor import HfqMAFactor
from factor.vol_ratio_factor import VolRatioFactor
from factor.rs_factor import RSFactor
from factor.main_index_factor import MainIndexFactor
from factor.vol_ma_factor import VolMAFactor
from factor.new_high_factor import NewHighFactor
from factor.median_factor import MedianFactor
from util.stock_util import get_code_trading_dates,judge_code_trading_date
"""
因子管理子系统
"""


class FactorModule:

    def get_single_stock_factors(self, code, factor,is_index, begin_date, end_date):
        """
        获取某只股票的某个因子在一段时间内的值
        :param code:  股票代码
        :param factor:  因子名称
        :param begin_date: 开始日期
        :param end_date: 结束日期
        :return: DataFrame(columns=['code',factor, 'date'])
        """

        factor_collection = DB_CONN[factor]
        actual_begin_date = judge_code_trading_date(code,is_index,begin_date)
        actual_end_date = judge_code_trading_date(code,is_index,end_date)
        factor_cursor = factor_collection.find(
            {'code': code, 'date': {'$gte': actual_begin_date, '$lte': actual_end_date},"index":is_index},
            sort=[('date', ASCENDING)],
            projection={'_id': False})
        trading_dates = get_code_trading_dates(code,is_index,begin_date, end_date)
        #print(trading_dates)
        factor_df = DataFrame([ x for x in factor_cursor if x['date'] in trading_dates])

        return factor_df

    def get_single_date_factors(self, factor, is_index,date):
        """
        获取某个因子在某个交易日的所有股票的数据
        :param factor: 因子名称
        :param date: 日期
        :return: DataFrame(columns=['code',factor, 'date'])
        """
        factor_collection = DB_CONN[factor]

        #没考虑指数的情况，存在问题
        factor_collection.create_index([('date', 1),('index',1)])
        trading_date = judge_code_trading_date(None,True,date)
        factor_cursor = factor_collection.find({'date': trading_date,"index":is_index},projection={'_id': False})

        factor_df = DataFrame(
            [x for x in factor_cursor])

        return factor_df

    def compute(self,begin_date=None,end_date=None):
        # 所有的因子实例
        factors = [
            # PEFactor(),
            ChangeRateFactor(),
            HfqMAFactor(),
            VolRatioFactor(),
            MainIndexFactor(),
            RSFactor(),
            VolMAFactor(),
            NewHighFactor(),
            MedianFactor()

        ]

        now = datetime.now().strftime('%Y-%m-%d')
        if begin_date is None:
            begin_date = "1995-02-24"
        if end_date is None:
            end_date = now
        for factor in factors:
            start_time = time.time()

            factor.compute(begin_date=begin_date, end_date=end_date)
            end_time = time.time()
            print('计算因子：%s, 从日期：%s 到 %s  ，共耗时：%.3f 秒' % (factor.name, begin_date, end_date, (end_time - start_time)), flush=True)
